//
//  SpeexCodeUtility.m
//  RainbowFM
//
//  Created by Kratos on 14/12/30.
//  Copyright (c) 2014年 RB. All rights reserved.
//

#import "ogg.h"
#import "os_types.h"
#import "speex.h"
#import "speex_bits.h"
#import "speex_buffer.h"
#import "speex_callbacks.h"
#import "speex_config_types.h"
#import "speex_echo.h"
#import "speex_header.h"
#import "speex_jitter.h"
#import "speex_preprocess.h"
#import "speex_stereo.h"
#import "speex_types.h"

#define SPX_FRAME_SZ        160 // PCM音频8khz*20ms -> 8000*0.02=160
#define SPX_MAX_NB_BYTES    200
#define SPX_SAMPLE_RATE     8000
#define SPX_QUALITY         3

typedef struct
{
    char chunkID[4];
    int  chunksize;
}ChunkHeader;

typedef struct
{
    short   formatTag;
    short   channels;
    int     samplesPerSec;
    int     avgBytesPerSec;
    short   blockAlign;
    short   bitsPerSample;
}WaveFormat;

typedef struct
{
    short   formatTag;
    short   channels;
    int     samplesPerSec;
    int     avgBytesPerSec;
    short   blockAlign;
    short   bitsPerSample;
    short   exsize;
}WaveFormatX;

typedef struct
{
    char    riffID[4];
    int     riffsize;
    char    riffFormat[4];
}RiffHeader;

typedef struct
{
    char        fmtID[4];
    int         fmtSize;
    WaveFormat  wf;
}FmtBlock;

uint16_t readUInt16(char* bis)
{
    uint16_t result = 0;
    result += ((uint16_t)(bis[0])) << 8;
    result += (uint8_t)(bis[1]);
    return result;
}

uint32_t readUInt32(char* bis)
{
    uint32_t result = 0;
    result += ((uint32_t) readUInt16(bis)) << 16;
    bis+=2;
    result += readUInt16(bis);
    return result;
}

int64_t readInt64(char* bis)
{
    int64_t result = 0;
    result += ((uint64_t) readUInt32(bis)) << 32;
    bis+=4;
    result += readUInt32(bis);
    return result;
}

static void writeWAVEHeader(NSMutableData* fpwave, int framesize)
{
    // 1. 写RIFF头
    RiffHeader riff ;
    
    memcpy(riff.riffID, "RIFF", 4);
    riff.riffsize = 4                   // WAVE
    + sizeof(ChunkHeader)               // fmt
    + sizeof(WaveFormatX)                // WaveFormatX
    + sizeof(ChunkHeader)               // DATA
    + framesize * 160 *sizeof(short);    //
    
    memcpy(riff.riffFormat, "WAVE", 4);
    
    [fpwave appendBytes:&riff length:sizeof(RiffHeader)];
    
    // 2. 写FMT块
    ChunkHeader chunk;
    WaveFormatX wfx;
    memcpy(chunk.chunkID, "fmt ", 4);
    chunk.chunksize = sizeof(WaveFormatX);
    [fpwave appendBytes:&chunk length:sizeof(ChunkHeader)];
    memset(&wfx, 0, sizeof(WaveFormatX));
    wfx.formatTag      = 1     ;
    wfx.channels       = 1     ; // 单声道
    wfx.samplesPerSec  = 8000  ; // 8khz
    wfx.avgBytesPerSec = 16000 ;
    wfx.blockAlign     = 2     ;
    wfx.bitsPerSample  = 16    ;  // 16位
    
    [fpwave appendBytes:&wfx length:sizeof(WaveFormatX)];
    
    // 3. 写data块头
    memcpy(chunk.chunkID, "data", 4);
    chunk.chunksize = framesize * 160 * sizeof(short);
    [fpwave appendBytes:&chunk length:sizeof(ChunkHeader)];
}


// 从WAVE文件读一个完整的PCM音频帧
// 返回值: 0-错误 >0: 完整帧大小
static int readPCMFrameData(short speech[], char* fpwave, int channels, int bitsPerSample)
{
    int nRead = 0;
    int x = 0, y=0;
    unsigned short ush1=0, ush2=0, ush=0;
    
    // 原始PCM音频帧数据
    uint8_t  pcmFrame_8b1[ SPX_FRAME_SZ ]       ;
    uint8_t  pcmFrame_8b2[ SPX_FRAME_SZ <<1]    ;
    uint16_t pcmFrame_16b1[ SPX_FRAME_SZ ]      ;
    uint16_t pcmFrame_16b2[ SPX_FRAME_SZ <<1]   ;
    
    nRead = (bitsPerSample / 8) *  SPX_FRAME_SZ  * channels;
    if (bitsPerSample ==8 && channels==1) {
        memcpy(pcmFrame_8b1,fpwave,nRead);
        for(x=0; x< SPX_FRAME_SZ ; x++) {
            speech[x] =(short)((short)pcmFrame_8b1[x] << 7);
        }
    } else if (bitsPerSample ==8 && channels == 2) {
        memcpy(pcmFrame_8b2,fpwave,nRead);
        for( x=0, y=0 ; y <  SPX_FRAME_SZ ; y++,x+=2 ) {
            ush1 = (short)pcmFrame_8b2[x+0];
            ush2 = (short)pcmFrame_8b2[x+1];
            ush = (ush1 + ush2) >> 1;
            speech[y] = (short)((short)ush << 7);
        }
    }
    else if (bitsPerSample == 16 &&  channels == 1) {
        memcpy(pcmFrame_16b1,fpwave,nRead);
        for(x=0; x <  SPX_FRAME_SZ ; x++) {
            speech[x] = (short)pcmFrame_16b1[x+0];
        }
    } else if (bitsPerSample ==16 && channels == 2) {
        memcpy(pcmFrame_16b2,fpwave,nRead);
        for( x=0, y=0; y <  SPX_FRAME_SZ ; y++,x+=2 ) {
            speech[y] = (short)((int)((int)pcmFrame_16b2[x+0] + (int)pcmFrame_16b2[x+1])) >> 1;
        }
    }
    
    // 如果读到的数据不是一个完整的PCM帧, 就返回0
    return nRead;
}

NSData *encodePCMToRawSpeex(char *PCMdata, int maxLen,int channels, int bitsPerSample)
{
    char *oldBuf = PCMdata;
    short speech[ SPX_FRAME_SZ ];
    int byte_counter, frames = 0, bytes = 0;
    
    spx_int16_t input[ SPX_FRAME_SZ ];
    char speexFrame[SPX_MAX_NB_BYTES];
    
    int tmp = SPX_QUALITY   ;
    void *encode_state = speex_encoder_init(&speex_nb_mode);
    speex_encoder_ctl(encode_state, SPEEX_SET_QUALITY, &tmp);
    
    SpeexPreprocessState *preprocess_state = speex_preprocess_state_init( SPX_FRAME_SZ , SPX_SAMPLE_RATE);
    
    int denoise = 1;
    int noiseSuppress = -10;
    speex_preprocess_ctl(preprocess_state, SPEEX_PREPROCESS_SET_DENOISE, &denoise);// 降噪
    speex_preprocess_ctl(preprocess_state, SPEEX_PREPROCESS_SET_NOISE_SUPPRESS, &noiseSuppress);// 噪音分贝数
    
    int agc = 1;
    int level = 24000;
    //actually default is 8000(0,32768),here make it louder for voice is not loudy enough by default.
    speex_preprocess_ctl(preprocess_state, SPEEX_PREPROCESS_SET_AGC, &agc);// 增益
    speex_preprocess_ctl(preprocess_state, SPEEX_PREPROCESS_SET_AGC_LEVEL,&level);// 增益后的值
    
    SpeexBits bits;
    speex_bits_init(&bits);
    NSMutableData *speexRawData = [[NSMutableData alloc] init];
    for (; ; ) {
        if ((PCMdata - oldBuf + sizeof(short)* SPX_FRAME_SZ ) > maxLen) {
            break;
        }
        
        speex_preprocess_run(preprocess_state, speech); //对原始录入声音进行处理(降噪、增益等)
        
        int nRead = readPCMFrameData(speech, PCMdata, channels, bitsPerSample);
        
        for (int i = 0; i <  SPX_FRAME_SZ ; i++) {
            input[i] = speech[i];
        }
        
        PCMdata += nRead;
        
        frames++;
        
        speex_bits_reset(&bits);
        speex_encode_int(encode_state, input, &bits);
        
        byte_counter = speex_bits_write(&bits, speexFrame, SPX_MAX_NB_BYTES);
        bytes += byte_counter;
        
        [speexRawData appendBytes:speexFrame length:byte_counter];
    }
    
    NSMutableData *speexData = [[NSMutableData alloc] init];
    SpeexHeader speexHeader;
    speex_init_header(&speexHeader, SPX_SAMPLE_RATE, 1, &speex_nb_mode);
    speexHeader.reserved1 = speex_bits_nbytes(&bits);
    [speexData appendBytes:&speexHeader length:speexHeader.header_size];
    [speexData appendData:speexRawData];
    
    speex_bits_destroy(&bits);
    speex_encoder_destroy(encode_state);
    speex_preprocess_state_destroy(preprocess_state);
    
    return speexData;
}

NSData *decodeSpeex2WAVE(NSData *data)
{
    if (data == nil){
        NSLog(@"data is nil...");
        return nil;
    }
    
    int nPos  = 0;
    char *buf = (char *)[data bytes];
    int maxLen = [data length];
    
    if (maxLen < sizeof(SpeexHeader)) {
        return nil;
    }
    
    SpeexHeader *speexHeader = (SpeexHeader *)buf;
    int nbBytes = speexHeader->reserved1;
    
    nPos += sizeof(SpeexHeader);
    if (nPos >= maxLen) {
        return nil;
    }
    
    //这时取出来的是纯speex数据
    buf += nPos;
    //--------------------------------------
    
    char *oldBuf = (char *)[data bytes];
    int frames = 0;
    
    short pcmFrame[ SPX_FRAME_SZ ];
    spx_int16_t output[ SPX_FRAME_SZ ];
    
    int tmp = 1;
    void *dec_state = speex_decoder_init(&speex_nb_mode);
    speex_decoder_ctl(dec_state, SPEEX_SET_ENH,&tmp);
    
    NSMutableData *PCMRawData = [[NSMutableData alloc] init];
    
    SpeexBits bits;
    speex_bits_init(&bits);
    for (; ; ) {
        if ((buf - oldBuf + nbBytes) > maxLen) {
            break;
        }
        
        speex_bits_read_from(&bits, buf, nbBytes);
        speex_decode_int(dec_state, &bits, output);
        
        for (int i = 0; i <  SPX_FRAME_SZ ; i++) {
            pcmFrame[i] = output[i];
        }
        
        [PCMRawData appendBytes:pcmFrame length:sizeof(short)* SPX_FRAME_SZ ];
        
        buf += nbBytes;
        frames++;
    }
    
    speex_bits_destroy(&bits);
    speex_decoder_destroy(dec_state);
    
    
    NSMutableData *outData = [[NSMutableData alloc]init];
    writeWAVEHeader(outData, frames);
    [outData appendData:PCMRawData];
    
    return outData;
}

float calculatePlayTime(NSData *speexData)
{
    const char* data = [speexData bytes];
    SpeexHeader* header = (SpeexHeader*)data;
    
    float play_time = 0.0;
    
    unsigned int rawSpeexDataLength = [speexData length] - header->header_size;
    play_time = (float)rawSpeexDataLength/(header->nb_channels * (header->frame_size) * 6 ); //每秒是50帧
    
    return play_time > 60.0 ? 60.0 : play_time;
}
